import glob
import os
import shutil
import matplotlib.pyplot as plt
import numpy as np
import tqdm
import sys

sys.path.append(os.path.dirname(os.path.realpath(__file__)))
# from punctuation import process_text
input_data_dir = '/home/work_nfs5_ssd/hfxue/data/data4w/source_1/huiting400'
output_data_dir = '/home/backup_nfs5/data_tts/huiting400'

py_source = '/home/backup_nfs5/data_tts/wenetspeech'
from gxl_ai_utils.utils import utils_file
from gxl_ai_utils.AiConstant import AI_LOGGER

logger = AI_LOGGER('./output/log/asr_handle_pipline.log')


def make_scp(root_dir: str, output_dir: str, file_name: str = 'wav.scp'):
    utils_file.makedir_sil(output_dir)
    res_dic = {}
    for root, dirs, files in os.walk(root_dir):
        for dirname in dirs:
            """"""
            dir_path = os.path.join(root, dirname)
            files = glob.glob(f"{dir_path}/*.wav")
            for file_path in tqdm.tqdm(files, total=len(files), desc=dirname):
                """"""
                filename = file_path.split("/")[-1]
                key = filename.split('.')[0]
                res_dic[key] = file_path
        break
    utils_file.write_dic_to_scp(res_dic, os.path.join(output_dir, file_name))
    print('over')


def make_mos_scp(root_dir: str, output_dir: str, file_name: str = 'mos.scp'):
    utils_file.makedir_sil(output_dir)
    res_dict = {}
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            # 检查文件扩展名是否为.scp
            if file.endswith(".scp"):
                # 构建完整的文件路径
                file_path = os.path.join(root, file)
                temp_dict = utils_file.load_dic_from_scp(file_path)
                res_dict.update(temp_dict)
    utils_file.write_dic_to_scp(res_dict, os.path.join(output_dir, file_name))


def show_image(dataset_name: str, input_list, output_dir: str):
    # 生成包含10000个数值的列表，范围为0-5之间
    utils_file.makedir_sil(output_dir)
    data = input_list
    total_lens = len(data)
    # 设置梯度为0.5
    bins = np.arange(0, 5.1, 0.5)
    # 绘制直方图
    plt.hist(data, bins=bins, edgecolor='black', alpha=0.7)

    # 标明每个柱子的具体数量
    for i in range(len(bins) - 1):
        count = np.sum((data >= bins[i]) & (data < bins[i + 1]))
        plt.text(bins[i] + 0.25, count + 50, str(count), ha='center', va='bottom')

    # 设置图表标题和标签
    plt.title(f'Histogram of {dataset_name}, total_lens:{total_lens}')
    plt.xlabel('Value')
    plt.ylabel('Frequency')

    # 显示图表
    plt.savefig(os.path.join(output_dir, 'histogram_{}.png'.format(dataset_name)))


def asr_handle_pipline(input_data_dir: str, output_data_dir: str):
    """"""

    text_path = os.path.join(input_data_dir, 'text')
    if not os.path.exists(text_path):
        logger.error('error: text file do not exist, return func')
        return
    wav_scp_path = os.path.join(input_data_dir, 'wav.scp')
    if not os.path.exists(wav_scp_path):
        logger.error('error: wav.scp file do not exist, return func')
        return
    dataset_name = os.path.basename(input_data_dir)

    logger.info(f'开始处理如下数据集:{dataset_name},得到wav_dict from wav.scp')
    # if os.path.exists(output_data_dir):
    #     shutil.rmtree(output_data_dir)
    os.makedirs(output_data_dir, exist_ok=True)
    if dataset_name == 'Wenetspeech':
        return
    wav_dict = utils_file.load_dic_from_scp(wav_scp_path)

    logger.info('得到原始的wav.scp')
    utils_file.write_dic_to_scp(wav_dict, os.path.join(output_data_dir, 'wav_origin.scp'))
    logger.info('得到原始的wav.scp完成')



    # 得到每个list文件包含100000条音频的目录
    logger.info('开始得到每个list文件包含100000条音频的temp目录， 为降噪做准备')
    temp_dir = os.path.join(output_data_dir, 'temp_lists_for_denoise')
    utils_file.makedir_sil(temp_dir)
    total_len = len(wav_dict)
    wav_path_list = list(wav_dict.values())
    temp_file_num = total_len // 100000 + 1
    for i in tqdm.tqdm(range(temp_file_num), desc='正在得到wav.list文件集合', total=temp_file_num):
        temp_file_path = os.path.join(temp_dir, f'wav_denoise_{i}')
        with open(temp_file_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(wav_path_list[i * 100000:(i + 1) * 100000]))

    logger.info('开始得到降噪音频，按照每个wav.list文件，得到不同的目录')
    return_code = os.system(f'bash ./shell_runner/denoise.sh {temp_dir} {output_data_dir}')
    if return_code == 0:
        logger.info('降噪完成')
        # 删除临时目录
        # shutil.rmtree(temp_dir)
        # logger.info('删除临时wav.list集合的目录及其所有文件')
    else:
        logger.error('降噪失败, 无法进行下一步处理， 退出')
        return False

    logger.info('开始得到降噪后的wav.scp')
    make_scp(os.path.join(output_data_dir, 'denoise'), output_data_dir, 'wav_denoise.scp')
    logger.info('wav.scp得到完毕')



    logger.info('开始得到每个scp文件包含100000条音频的temp目录，为降采样做准备')
    temp_dir = os.path.join(output_data_dir, 'temp_scps_for_ds')
    utils_file.makedir_sil(temp_dir)
    wav_denoise_scp_path = os.path.join(output_data_dir, 'wav_denoise.scp')
    wav_denoise_dict = utils_file.load_dic_from_scp(wav_denoise_scp_path)
    total_len = len(wav_denoise_dict)
    items_list = list(wav_denoise_dict.items())
    temp_file_num = total_len // 100000 + 1
    for i in tqdm.tqdm(range(temp_file_num), desc='正在得到wav.scp文件集合', total=temp_file_num):
        temp_file_path = os.path.join(temp_dir, f'wav_ds_{i}')
        item_list_part = items_list[i * 100000:(i + 1) * 100000]
        dict_temp = dict(item_list_part)
        utils_file.write_dic_to_scp(dict_temp, temp_file_path)
    logger.info('temp目录得到完毕')

    logger.info('开始降采样')
    return_code = os.system(f'bash ./shell_runner/ds.sh {temp_dir} {output_data_dir}')
    if return_code == 0:
        logger.info('降采样完成')
        # 删除临时目录
        # shutil.rmtree(temp_dir)
        # logger.info('删除临时wav.scp集合的目录及其所有文件')
    else:
        logger.error('降采样失败, 无法进行下一步处理， 退出')
        return False

    logger.info('开始得到降采样后的wav.scp')
    make_scp(os.path.join(output_data_dir, 'dswavs'), output_data_dir, 'wav_ds.scp')
    logger.info('wav.scp得到完毕')



    logger.info('开始得到每个scp文件包含100000条音频的temp目录，为mos做准备')
    temp_dir = os.path.join(output_data_dir, 'temp_scps_for_mos')
    utils_file.makedir_sil(temp_dir)
    wav_ds_scp_path = os.path.join(output_data_dir, 'wav_ds.scp')
    wav_ds_dict = utils_file.load_dic_from_scp(wav_ds_scp_path)
    total_len = len(wav_ds_dict)
    items_list = list(wav_ds_dict.items())
    temp_file_num = total_len // 100000 + 1
    for i in tqdm.tqdm(range(temp_file_num), desc='正在得到wav.scp文件集合', total=temp_file_num):
        temp_file_path = os.path.join(temp_dir, f'wav_mos_{i}')
        item_list_part = items_list[i * 100000:(i + 1) * 100000]
        dict_temp = dict(item_list_part)
        utils_file.write_dic_to_scp(dict_temp, temp_file_path)
    logger.info('temp目录得到完毕')

    logger.info('开始mos打分')
    return_code = os.system(f'bash ./shell_runner/mos.sh {temp_dir} {output_data_dir}')
    if return_code == 0:
        logger.info('降采样完成')
        # 删除临时目录
        # shutil.rmtree(temp_dir)
        # logger.info('删除临时wav.scp集合的目录及其所有文件')
    else:
        logger.error('降采样失败, 无法进行下一步处理， 退出')
        return False

    logger.info('开始得到mos.scp')
    make_mos_scp(os.path.join(output_data_dir, 'mos'), output_data_dir, 'mos.scp')
    logger.info('mos.scp得到完毕')




    logger.info('开始绘制直方图')
    mos_dict = utils_file.load_dic_from_scp(os.path.join(output_data_dir, 'mos.scp'))
    mos_num_list = list(mos_dict.values())
    float_array = np.array(mos_num_list, dtype=float)
    show_image(dataset_name, float_array, output_data_dir)
    logger.info('直方图绘制完毕')

    # logger.info('开始处理text')
    # process_text(text_path, output_data_dir)
    # logger.info('text处理完毕')


if __name__ == '__main__':
    asr_handle_pipline(input_data_dir, output_data_dir)
